package cn.fetosoft.woodpecker.core.jmeter;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.Map;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

/**
 * 动态加载扩展jar包
 * @author guobingbing
 * @version 1.0
 * @wechat t_gbinb
 * @create 2021/8/7 20:48
 */
public final class DynamicExtensionLoader extends URLClassLoader {

	private static final Logger log = LoggerFactory.getLogger(DynamicExtensionLoader.class);
	//保存本类加载器加载的class
	private static final Map<String, byte[]> BYTE_CODE_MAP = new HashMap<>();

	public DynamicExtensionLoader(URL[] urls) {
		super(urls);
		this.init(urls);
	}

	private void init(URL[] urls){
		try{
			for(URL url : urls) {
				this.loadClasses(new JarFile(url.getPath()));
			}
		}catch(Exception e){
			log.error("init", e);
		}
	}

	@Override
	public Class<?> findClass(final String name) throws ClassNotFoundException{
		if(BYTE_CODE_MAP.containsKey(name)){
			byte[] byteCode = BYTE_CODE_MAP.get(name);
			return this.defineClass(name, byteCode, 0, byteCode.length);
		}else{
			return super.findClass(name);
		}
	}

	private void loadClasses(JarFile jarFile){
		Enumeration<JarEntry> en = jarFile.entries();
		InputStream input = null;
		try{
			while (en.hasMoreElements()) {
				JarEntry je = en.nextElement();
				String name = je.getName();
				//这里添加了路径扫描限制
				if (name.endsWith(".class")) {
					String className = name.replace(".class", "").replaceAll("/", ".");
					input = jarFile.getInputStream(je);
					ByteArrayOutputStream baos = new ByteArrayOutputStream();
					int bufferSize = 4096;
					byte[] buffer = new byte[bufferSize];
					int bytesNumRead = 0;
					while ((bytesNumRead = input.read(buffer)) != -1) {
						baos.write(buffer, 0, bytesNumRead);
					}
					byte[] classBytes = baos.toByteArray();
					BYTE_CODE_MAP.put(className, classBytes);
				}
			}
		} catch (IOException e) {
			log.error("loadClasses", e);
		} finally {
			if(input!=null){
				try {
					input.close();
				} catch (IOException e) {
					e.printStackTrace();
				}
			}
		}

		for (Map.Entry<String, byte[]> entry : BYTE_CODE_MAP.entrySet()) {
			String key = entry.getKey();
			try {
				//载入所有的class
				this.loadClass(key);
			} catch (ClassNotFoundException e) {
				e.printStackTrace();
			} catch (NoClassDefFoundError e){
				//忽略jar中的三方依赖缺少的异常
				System.out.println("NoClassDefFoundError =>" + key);
			}
		}
	}

	public static void main(String[] args) {
		try {
			//加载扩展jar包
			URL[] urls = new URL[]{new File("D:\\Git\\meterweb\\meterweb-web\\lib\\mysql-connector-java-5.1.47.jar").toURI().toURL()};
			DynamicExtensionLoader loader = new DynamicExtensionLoader(urls);

			Class<?> driverClazz = loader.loadClass("com.mysql.jdbc.Driver");
			System.out.println("DataSourceElement =>" + driverClazz.getDeclaredConstructor().newInstance());
			System.out.println(driverClazz.getClassLoader().toString());
		}catch (Exception e){
			e.printStackTrace();
		}
	}
}
